#include "model_yolo_v6.h"                                  // 包含 YOLO v6 类定义
//#include "../txr_algo_dlm_v5/common/common.hpp"       // 包含常用定义和工具函数
#include "yaml-cpp/yaml.h"                            // YAML 解析库，用于读取配置文件
#include "boost/process.hpp"                          // Boost 进程库，用于进程管理
#include "common/device_utility.h"
#include "boost/filesystem/operations.hpp"            // Boost 文件系统库，用于文件和目录操作
#include "boost/filesystem/path.hpp"
#include "boost/filesystem/fstream.hpp"

//#include "dll.h"                                      // DLL 库（具体功能未在代码中显示）

using namespace CryptoPP;                             // 使用 Crypto++ 命名空间



std::vector<char> decryptENNX(
	const std::vector<char> & v_encrypt,
	const std::string& key,
	const std::string& iv) {
	std::vector<char> v_decrypt;                            // 存储解密结果
	try {
		// 初始化 CBC 模式的 AES 解密对象
		CBC_Mode<AES>::Decryption decryption(reinterpret_cast<const byte *>(key.data()), key.size(), (const byte*)iv.data());

		std::vector<char> _decrypt = v_encrypt;             // 复制加密数据
		decryption.ProcessData(
			reinterpret_cast<CryptoPP::byte*>(_decrypt.data()), // 解密后的数据存储位置
			reinterpret_cast<const CryptoPP::byte*>(v_encrypt.data()), // 原始加密数据
			v_encrypt.size());

		// 处理填充，去除解密数据末尾的填充字节
		byte paddingByte = _decrypt.back();                  // 获取填充字节
		size_t paddingSize = static_cast<size_t>(paddingByte);
		_decrypt.resize(v_encrypt.size() - paddingSize);     // 删除填充字节

		v_decrypt = _decrypt;                                // 将解密数据存储到返回变量
	} catch (const CryptoPP::Exception& e) {
		std::cerr << e.what() << std::endl;                 // 捕获并输出解密异常
	}
	return v_decrypt;                                       // 返回解密数据
}



/**
 * @brief 读取文件数据并返回字符向量
 * @param file_path 文件路径
 * @return 返回文件内容的字符向量
 */
// 读取文件内容到 vector<char>
std::vector<char> get_file_data(const std::string& file_path) {
	std::ifstream file(file_path, std::ios_base::binary);
	if (!file.is_open()) {
		throw std::runtime_error("Failed to open file: " + file_path);
	}
	file.seekg(0, std::ios_base::end);
	const size_t file_len = file.tellg();
	file.seekg(0, std::ios_base::beg);
	std::vector<char> data(file_len, 0);
	if (file.read(data.data(), file_len)) {
		return data;
	} else {
		throw std::runtime_error("Failed to read file: " + file_path);
	}
}
/**
 * @brief 将字符向量数据写入文件
 * @param data 要写入的数据
 * @param file_path 文件路径
 */
void write_file_data(std::vector<char> data, std::string file_path) {
	std::ofstream file(file_path, std::ios_base::binary);   // 以二进制方式打开文件写入
	file.write(data.data(), data.size());                   // 写入向量数据到文件
	file.close();                                           // 关闭文件
}

// 对数据进行简单的加密操作。
// @param data 待加密的数据。
void encrypt_data(std::vector<char>& data) {
	char key = 0x1; // 简单的密钥，实际情况下应更复杂
	for (size_t i = 0; i < data.size(); ++i) {
		data[i] ^= key; // 异或操作进行加密
	}
}

// 对数据进行简单的解密操作。
// @param data 待解密的数据。
void decrypt_data(std::vector<char>& data) {
	char key = 0x1; // 与加密相同的密钥
	for (size_t i = 0; i < data.size(); ++i) {
		data[i] ^= key; // 异或操作进行解密
	}
}

void merge_files(
    const std::string& xml_path,
    const std::string& bin_path,
    const st_encrypt_info_ovino& info,
    const std::string& output_path) {
    std::vector<char> xml_data = get_file_data(xml_path);
    std::vector<char> bin_data = get_file_data(bin_path);

    // 打印解析后的XML和BIN数据的前20个字符
    #ifdef TEST_ENCRYPTION
    std::cout << "XML Data (first 20 bytes): ";
    for (size_t i = 0; i < 20; ++i) {
        std::cout << std::hex << (int)xml_data[i] << " ";
    }
    std::cout << std::endl;

    std::cout << "BIN Data (first 20 bytes): ";
    for (size_t i = 0; i < 20; ++i) {
        std::cout << std::hex << (int)bin_data[i] << " ";
    }
    std::cout << std::endl;
    #endif

    // 加密XML和BIN数据
    encrypt_data(xml_data);
    encrypt_data(bin_data);

    // 创建头部信息
    st_encrypt_header header = {
        sizeof(st_encrypt_info_ovino), // 结构体大小
        xml_data.size(),               // XML文件大小
        bin_data.size()                // BIN文件大小
    };

    // 将结构体和文件数据合并到一个 vector 中
    std::vector<char> merged_data;
    merged_data.reserve(sizeof(st_encrypt_header) + sizeof(st_encrypt_info_ovino) + xml_data.size() + bin_data.size());
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&header), reinterpret_cast<const char*>(&header) + sizeof(st_encrypt_header));
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&info), reinterpret_cast<const char*>(&info) + sizeof(st_encrypt_info_ovino));
    merged_data.insert(merged_data.end(), xml_data.begin(), xml_data.end());
    merged_data.insert(merged_data.end(), bin_data.begin(), bin_data.end());

    // 写入文件
    std::ofstream output_file(output_path, std::ios_base::binary);
    if (!output_file.is_open()) {
        throw std::runtime_error("Failed to open output file: " + output_path);
    }
    output_file.write(merged_data.data(), merged_data.size());
    if (!output_file) {
        throw std::runtime_error("Failed to write to output file: " + output_path);
    }
}
/**
 * @brief 从 YAML 配置文件读取数据到 st_encrypt_info 结构体
 * @param input 转换输入结构，包含配置文件路径和其他信息
 * @return 填充后的 st_encrypt_info 结构体
 * @throws std::runtime_error 如果配置文件有问题或读取失败
 * dynamic没有读取！！！！
 */
st_encrypt_info_ovino loadConfigToInfo(const st_trans_model_input& input) {
    // 加载 YAML 配置文件
    YAML::Node root = YAML::LoadFile(input.cfg_path); // 加载 YAML 配置文件
    YAML::Node config;

    // 检查并读取 YOLOv6 或 DET_LB 节点
    if (root["YOLOv6"].IsDefined()) {
        config = root["YOLOv6"]; // 如果存在 YOLOv6 节点，读取配置
    } else if (root["DET_LB"].IsDefined()) {
        config = root["DET_LB"]; // 如果存在 DET_LB 节点，读取配置
    } else {
        throw std::runtime_error("Invalid configuration: missing YOLOv6 or DET_LB node.");
    }

    // 初始化结构体，存储配置信息
    st_encrypt_info_ovino info;

    // 填充批次大小
    if (!config["BATCH_SIZE"].IsDefined()) {
        throw std::runtime_error("BATCH_SIZE undefined.");
    }
    info.batch_size = config["BATCH_SIZE"].as<int>();

    // 检查并填充分类数量
    if (!config["CLASS_NUM"].IsDefined()) {
        throw std::runtime_error("CLASS_NUM undefined.");
    }
    info.class_num = config["CLASS_NUM"].as<int>();

    // 填充输入通道数、图像宽高
    if (!config["INPUT_CHANNEL"].IsDefined() ||
        !config["IMAGE_WIDTH"].IsDefined() ||
        !config["IMAGE_HEIGHT"].IsDefined()) {
        throw std::runtime_error("Input channel or image size undefined.");
    }
    info.input_channel = config["INPUT_CHANNEL"].as<int>();
    info.image_width = config["IMAGE_WIDTH"].as<int>();
    info.image_height = config["IMAGE_HEIGHT"].as<int>();

    // 填充目标检测阈值和非极大值抑制阈值
    if (!config["obj_threshold"].IsDefined() ||
        !config["nms_threshold"].IsDefined()) {
        throw std::runtime_error("Detection thresholds undefined.");
    }
    info.obj_threshold = config["obj_threshold"].as<float>();
    info.nms_threshold = config["nms_threshold"].as<float>();

    // 读取锚点和步长数据
    if (!config["num_anchors"].IsDefined() ||
        !config["strides"].IsDefined()) {
        throw std::runtime_error("Anchors or strides undefined.");
    }

	// 锚点和步长数组
	info.num_anchors_size = root["DET_LB"]["num_anchors"].size();
	for (size_t i = 0; i < info.num_anchors_size; ++i) {
		info.num_anchors[i] = root["DET_LB"]["num_anchors"][i].as<int>();
	}
	info.stride_num = root["DET_LB"]["strides"].size();
	for (size_t i = 0; i < info.stride_num; ++i) {
		info.strides[i] = root["DET_LB"]["strides"][i].as<int>();
	}

    return info; // 返回填充完成的结构体
}

/**
 * @brief 读取模型文件并解密（支持 .onnx 和 .ennx 文件）
 * @param input 转换输入结构，包含模型路径等信息
 * @param key 解密密钥
 * @param iv 解密初始化向量
 * @return 模型文件数据的字节向量
 * @throws std::runtime_error 如果文件格式不支持或读取失败
 */
std::vector<char> loadAndDecrypt_ONNX_ENNX_Model(
	const st_trans_model_input& input,
	const std::string& key,
	const std::string& iv)
{
	namespace fs = boost::filesystem;

	// 检查模型文件的扩展名
	fs::path p(input.model_path.c_str());
	if (p.extension().string() == ".onnx") {
		// 如果是 `.onnx` 模型文件，直接读取数据
		return get_file_data(p.string());
	} else if (p.extension().string() == ".ennx") {
		// 如果是加密的 `.ennx` 模型文件
		std::vector<char> v_e_model = get_file_data(p.string()); // 加载加密模型文件
		return decryptENNX(v_e_model, key, iv);                 // 使用密钥和向量解密数据
	} else {
		// 文件格式不支持
		throw std::runtime_error("Unsupported model file format: " + p.extension().string());
	}
}

void remove_temp_model_files(const std::string& directory = ".") {
	for (const auto& entry : std::filesystem::directory_iterator(directory)) {
		if (entry.is_regular_file()) {
			std::string filename = entry.path().filename().string();
			if (filename.rfind("temp_model.", 0) == 0) { // 检查文件名是否以 "temp_model." 开头
				try {
					std::filesystem::remove(entry.path());
					// std::cout << "Deleted: " << entry.path() << std::endl;
				} catch (const std::filesystem::filesystem_error& e) {
					std::cerr << "Error deleting " << entry.path() << ": " << e.what() << std::endl;
				}
			}
		}
	}
}


/**
 * @brief 将 ONNX 模型数据转换为 IR 格式的 XML 和 BIN 文件数据缓冲区
 * @param core OpenVINO 核心对象
 * @param model_data ONNX 模型数据的字符串表示
 * @return 一个包含 XML 和 BIN 文件数据的缓冲区对
 */
std::pair<std::vector<char>,std::vector<char>>
coreModel_to_XB_buffer(
	ov::Core& core,
	const std::string& model_data)
{
	// 加载模型
	//todo 这里有个库的bug，会产生一个kernel.errors.txt但是不影响任何推理
	std::shared_ptr<ov::Model> model = core.read_model(model_data, ov::Tensor{});

	// 临时文件路径
	const std::string temp_xml_path = "temp_model.xml";
	const std::string temp_bin_path = "temp_model.bin";

	// 将调整后的模型序列化到临时文件
	ov::serialize(model, temp_xml_path, temp_bin_path);

	// 读取临时文件内容到内存
	std::ifstream xml_file(temp_xml_path, std::ios::binary);
	std::ifstream bin_file(temp_bin_path, std::ios::binary);

	std::vector<char> xml_buffer((std::istreambuf_iterator<char>(xml_file)), std::istreambuf_iterator<char>());
	std::vector<char> bin_buffer((std::istreambuf_iterator<char>(bin_file)), std::istreambuf_iterator<char>());

	// 打印文件大小
	std::cout << "XML data size: " << xml_buffer.size() << " bytes" << std::endl;
	std::cout << "BIN data size: " << bin_buffer.size() << " bytes" << std::endl;

	// 清理临时文件
	std::remove(temp_xml_path.c_str());
	std::remove(temp_bin_path.c_str());

	// 返回 XML 和 BIN 缓冲区
	return {xml_buffer, bin_buffer};
}


void encrypt_XML_BIN_to_hdats(
    std::vector<char>& xml_data,
    std::vector<char>& bin_data,
    const st_encrypt_info_ovino& info,
    const std::string& output_path)
{

    std::cout << std::endl;
    // 加密XML和BIN数据
    encrypt_data(xml_data);
    encrypt_data(bin_data);

    // 创建头部信息
    st_encrypt_header header = {
        sizeof(st_encrypt_info_ovino), // 结构体大小
        xml_data.size(),               // XML文件大小
        bin_data.size()                // BIN文件大小
    };

    // 将结构体和文件数据合并到一个 vector 中
    std::vector<char> merged_data;
    merged_data.reserve(sizeof(st_encrypt_header) + sizeof(st_encrypt_info_ovino) + xml_data.size() + bin_data.size());
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&header), reinterpret_cast<const char*>(&header) + sizeof(st_encrypt_header));
    merged_data.insert(merged_data.end(), reinterpret_cast<const char*>(&info), reinterpret_cast<const char*>(&info) + sizeof(st_encrypt_info_ovino));
    merged_data.insert(merged_data.end(), xml_data.begin(), xml_data.end());
    merged_data.insert(merged_data.end(), bin_data.begin(), bin_data.end());

    // 写入文件
    std::ofstream output_file(output_path, std::ios_base::binary);
    if (!output_file.is_open()) {
        throw std::runtime_error("Failed to open output file: " + output_path);
    }
    output_file.write(merged_data.data(), merged_data.size());
    if (!output_file) {
        throw std::runtime_error("Failed to write to output file: " + output_path);
    }
}

/**
 * @brief YOLO v6 构造函数
 */
yolo_v6::yolo_v6() = default;

/**
 * @brief YOLO v6 析构函数
 */
yolo_v6::~yolo_v6() = default;

/**
 * @brief 转换 YOLO v6 模型文件
 * @param input 转换输入结构，包含配置文件路径和其他信息
 * @return 转换成功返回 true，否则返回 false
 */
bool yolo_v6::TransModel(const st_trans_model_input& input)
{
	// 定义加密用的模型参数信息
	st_encrypt_info_ovino info;
	// 定义ONNX模型缓存向量
	std::vector<char> v_model;

	try {
		info= loadConfigToInfo(input);
		// 后续代码逻辑可以使用 info
	} catch (const std::runtime_error& e) {
		std::cerr << "Error loading configuration: " << e.what() << std::endl;
		return false;
	}
	try {
		const std::string key = "irkdh_573%3?iq5h";   // 解密密钥
		const std::string iv = "1234567890123456";    // 初始化向量
		v_model = loadAndDecrypt_ONNX_ENNX_Model(input, key, iv);
		// 后续代码逻辑可以使用 v_model
	}
	catch (const std::runtime_error& e) {
		std::cerr << "Error loading model: " << e.what() << std::endl;
		return false;
	}

	// todo 将 ONNX 模型转换为 OpenVINO IR引擎模型
	// 初始化 OpenVINO 核心对象
	ov::Core core;
	// 将vector<char>转为read_model的输入形式string
	const std::string model_data(v_model.begin(), v_model.end());

	// 调用函数转换模型为缓冲区
	const auto [xml_, bin_] =
		coreModel_to_XB_buffer(core, model_data);

	// 分别获取 XML 和 BIN 数据缓冲区
	std::vector<char> xml_buffer = xml_;
	std::vector<char> bin_buffer = bin_;

	//todo 加密模型和模型信息（2）
	// 填充引擎数据相关信息 2
	info.engine_size = xml_buffer.size() + bin_buffer.size();
	info.gpu_model[0] = *"Intel core";          // 初始化 GPU 型号字符串

	// 获取加密文件的保存路径
	const std::string encrypt_file = input.trans_path;
	// 调用加密函数
	encrypt_XML_BIN_to_hdats(
		xml_buffer,
		bin_buffer,
		info,
		encrypt_file);

	remove_temp_model_files();
	// 转换成功，返回 true
	return true;
}
